{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "863c8a56-856d-454a-864e-2ece5ea1e7c7",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/conda/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "import xarray as xr\n",
    "import numpy as np\n",
    "import tqdm\n",
    "from scipy.stats import pearsonr\n",
    "from skimage.measure import block_reduce\n",
    "from src.configuration import Config\n",
    "import src.utils.xarray_utils as xu"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "9d6b312c-57a1-44ae-8fff-bdf5192fc9ac",
   "metadata": {},
   "outputs": [],
   "source": [
    "# load datasets\n",
    "scratch_path = '/dss/dssfs04/lwp-dss-0002/pn49fu/pn49fu-dss-0002/ge45tac2/diffusion-downscaling/'\n",
    "\n",
    "config = Config(out_path=scratch_path)\n",
    "\n",
    "test_split = slice('2004', '2018')\n",
    "train_split = slice('1950', '1990')\n",
    "\n",
    "config.target_filename = 'era5_4x_downscaled_v3.nc'\n",
    "era5_train = xr.open_dataset(config.data_path + '/' + config.target_filename, chunks={'time': 1}).precipitation.sel(time=train_split)*(24*3600)\n",
    "era5 = xr.open_dataset(config.data_path + '/' + config.target_filename).precipitation.sel(time=test_split)*(24*3600)\n",
    "\n",
    "config.esm_filename = 'poem_historical_3deg_4xdown_bilinear_lowpass_dqm.nc'\n",
    "esm_hr_train = xr.open_dataset(config.data_path + '/' + config.esm_filename, chunks={'time': 1}).precipitation.sel(time=train_split)*(24*3600)\n",
    "esm_hr = xr.open_dataset(config.data_path + '/' + config.esm_filename).precipitation.sel(time=test_split)*(24*3600)\n",
    "esm_hr = xu.shift_longitudes(esm_hr)\n",
    "\n",
    "config.esm_filename = 'poem_historical_3deg_4xdown_bilinear_v3.nc'\n",
    "esm_hr_no_qm = xr.open_dataset(config.data_path + '/' + config.esm_filename).precipitation.sel(time=test_split)*(24*3600)\n",
    "esm_hr_no_qm = xu.shift_longitudes(esm_hr_no_qm)\n",
    "\n",
    "esm_lr_filename = 'poem_historical_3deg.nc'\n",
    "esm_lr = xr.open_dataset(config.data_path + '/' + esm_lr_filename, chunks={'time': 1}).precipitation.sel(time=test_split)*(24*3600)\n",
    "esm_lr = xu.shift_longitudes(esm_lr)\n",
    "\n",
    "sde_filename = scratch_path+'results/sde_historical_poem_4x_dqm_2004_2018.nc'\n",
    "sde = xr.open_dataset(sde_filename, chunks={'time': 1}).precipitation.sel(time=test_split)*(24*3600)\n",
    "sde = xu.shift_longitudes(sde)\n",
    "\n",
    "cm_filename = scratch_path+'results/cm_historical_scale_468.nc'\n",
    "cm = xr.open_dataset(cm_filename, chunks={'time': 1}).precipitation*(24*3600)\n",
    "cm = xu.shift_longitudes(cm)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "2f7e1695-eb95-453c-be31-35c3c2f1c4d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# helper functions\n",
    "def pool(x, stride):\n",
    "    return block_reduce(x, (stride,stride), np.mean)\n",
    "\n",
    "def correlation(a, b):\n",
    "    return pearsonr(a.flatten(), b.flatten()).statistic"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c46ba1b6-d7c1-4c4c-85c6-7ea538c9ae63",
   "metadata": {},
   "source": [
    "# Table"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "id": "fae239a3-e884-42c5-869d-c98ecd51907e",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 5475/5475 [01:57<00:00, 46.63it/s]\n"
     ]
    }
   ],
   "source": [
    "num_t = len(sde)\n",
    "sde_corr_pooled = np.zeros(num_t)\n",
    "cm_corr_pooled = np.zeros(num_t)\n",
    "\n",
    "for t in tqdm.tqdm(range(num_t)):\n",
    "    esm_slice = esm_hr.isel(time=t)\n",
    "    esm_pooled = pool(esm_slice.values, stride=4)\n",
    "\n",
    "    sde_slice = sde.isel(time=t)\n",
    "    sde_pooled = pool(sde_slice.values, stride=4)\n",
    "    sde_corr_pooled[t] = correlation(sde_pooled, esm_pooled)\n",
    "    \n",
    "    cm_slice = cm.isel(time=t)\n",
    "    cm_pooled = pool(cm_slice.values, stride=4)\n",
    "    cm_corr_pooled[t] = correlation(cm_pooled, esm_pooled)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 73,
   "id": "4d0a4580-1836-47d7-b32c-7925e5efb736",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SDE correlation: mean = 0.918, std = 0.014\n",
      "CM correlation: mean = 0.954, std = 0.008\n"
     ]
    }
   ],
   "source": [
    "print(f\"SDE correlation: mean = {sde_corr_pooled.mean():2.3f}, std = {sde_corr_pooled.std():2.3f}\")\n",
    "print(f\"CM correlation: mean = {cm_corr_pooled.mean():2.3f}, std = {cm_corr_pooled.std():2.3f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 73,
   "id": "e35953b6-50e6-4e85-a64c-5eacef696975",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 5475/5475 [00:41<00:00, 131.85it/s]\n"
     ]
    }
   ],
   "source": [
    "from src.utils import compute_filtered_correlations, low_pass_filter\n",
    "\n",
    "sde_lp = low_pass_filter(sde, cutoff=0.0667)\n",
    "cm_lp = low_pass_filter(cm, cutoff=0.0667)\n",
    "\n",
    "num_t = len(sde)\n",
    "sde_corr_filtered = np.zeros(num_t)\n",
    "cm_corr_filtered = np.zeros(num_t)\n",
    "\n",
    "for t in tqdm.tqdm(range(num_t)):\n",
    "    esm_slice = esm_hr[t].values\n",
    "    \n",
    "    sde_slice = sde_lp[t]\n",
    "    sde_corr_filtered[t] = correlation(sde_slice, esm_slice)\n",
    "    \n",
    "    cm_slice = cm_lp[t]\n",
    "    cm_corr_filtered[t] = correlation(cm_slice, esm_slice)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "id": "2164a1fc-88cd-441a-9687-cf6b7021543f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SDE correlation: mean = 0.916, std = 0.012\n",
      "CM correlation: mean = 0.941, std = 0.009\n"
     ]
    }
   ],
   "source": [
    "print(f\"SDE correlation: mean = {sde_corr_filtered.mean():2.3f}, std = {sde_corr_filtered.std():2.3f}\")\n",
    "print(f\"CM correlation: mean = {cm_corr_filtered.mean():2.3f}, std = {cm_corr_filtered.std():2.3f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "5183a488-d8a0-463e-a0bb-155fc5bd875c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SDE MAE = 0.214\n",
      "CM MAE = 0.217\n",
      "ESM MAE = 0.778\n"
     ]
    }
   ],
   "source": [
    "esm_hr_no_qm['time'] =  cm.time \n",
    "cm_mae = abs(era5.mean(dim=\"time\") - cm.mean(dim=\"time\")).mean().values\n",
    "sde_mae = abs(era5.mean(dim=\"time\") - sde.mean(dim=\"time\")).mean().values\n",
    "esm_mae = abs(era5.mean(dim=\"time\") - esm_hr_no_qm.mean(dim=\"time\")).mean().values\n",
    "\n",
    "print(f\"SDE MAE = {sde_mae:2.3f}\")\n",
    "print(f\"CM MAE = {cm_mae:2.3f}\")\n",
    "print(f\"ESM MAE = {esm_mae:2.3f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "641050eb-bcc2-42f7-a651-908fec6845fb",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SDE rel. error reduction = -72.506\n",
      "CM rel. error reduction = -72.081\n"
     ]
    }
   ],
   "source": [
    "def relative_error_reduction(error_1, error_2):\n",
    "    return (error_1 - error_2) / error_2 * 100.0\n",
    "    \n",
    "print(f\"SDE rel. error reduction = {relative_error_reduction(sde_mae, esm_mae):2.3f}\")\n",
    "print(f\"CM rel. error reduction = {relative_error_reduction(cm_mae, esm_mae):2.3f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "7ecbeb57-8850-4eb2-90d0-7866d6e44b88",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SDE MAE = 1.106\n",
      "CM MAE = 1.080\n",
      "ESM MAE = 3.474\n"
     ]
    }
   ],
   "source": [
    "esm_hr_no_qm['time'] =  cm.time \n",
    "q = 0.95\n",
    "cm_mae = abs(era5.load().quantile(q, dim=\"time\") - cm.load().quantile(q, dim=\"time\")).mean().values\n",
    "sde_mae = abs(era5.load().quantile(q, dim=\"time\") - sde.load().quantile(q, dim=\"time\")).mean().values\n",
    "esm_mae = abs(era5.load().quantile(q, dim=\"time\") - esm_hr_no_qm.load().quantile(q, dim=\"time\")).mean().values\n",
    "\n",
    "print(f\"SDE MAE = {sde_mae:2.3f}\")\n",
    "print(f\"CM MAE = {cm_mae:2.3f}\")\n",
    "print(f\"ESM MAE = {esm_mae:2.3f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "ac1a8ff1-22a5-444e-a032-982da3887271",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SDE rel. error reduction = -68.154\n",
      "CM rel. error reduction = -68.920\n"
     ]
    }
   ],
   "source": [
    "def relative_error_reduction(error_1, error_2):\n",
    "    return (error_1 - error_2) / error_2 * 100.0\n",
    "    \n",
    "print(f\"SDE rel. error reduction = {relative_error_reduction(sde_mae, esm_mae):2.3f}\")\n",
    "print(f\"CM rel. error reduction = {relative_error_reduction(cm_mae, esm_mae):2.3f}\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
